# coding: utf-8

import torch
from torch import nn
import numpy as np
import os
import time


__all__ = [
    "awgn_torch", "rayleigh_torch", "rician_torch", "Channel"
]


def awgn_torch(input_signal, snr_dB, rate=1.):
    print(f"输入设置的信噪比（dB）： {snr_dB} dB")
    device = input_signal.device
    # N = input_signal.shape[0]
    # avg_energy = (torch.abs(input_signal) ** 2).sum() / N
    avg_energy = torch.mean(torch.abs(input_signal) ** 2, dtype=torch.float32)
    snr_linear = 10 ** (snr_dB / 10.)
    noise_variance = avg_energy / (2 * rate * snr_linear)

    if torch.is_complex(input_signal):
        noise = (torch.randn_like(input_signal, dtype=torch.float32, device=device) +
                 1j * torch.randn_like(input_signal, dtype=torch.float32, device=device)) \
                * noise_variance.sqrt()
    else:
        noise = torch.randn_like(input_signal, dtype=torch.float32, device=device) * (2 * noise_variance).sqrt()

    output_signal = input_signal + noise
    ps, pn = torch.mean(torch.abs(input_signal)**2, dtype=torch.float32), torch.mean(torch.abs(noise)**2, dtype=torch.float32)
    print(f"信号输出时>> 信号功率(ps)：{ps}, 噪声功率(pn): {pn}, 信噪比（10*lg(ps/pn)）：{10*torch.log10(ps / pn)} dB")
    return output_signal


def rayleigh_torch(input_signal, snr_dB):
    N = input_signal.shape[0]
    device = input_signal.device
    # 生成 rayleigh 信道系数 h
    h = (torch.randn(N, device=device) +1j * torch.randn(N, device=device)) / torch.sqrt(torch.tensor(2., device=device))
    # 信道作用
    # bits_times_channel_coefficients = input_signal * h
    # 加噪声
    bits_times_channel_coefficients_with_noise = awgn_torch(input_signal, snr_dB)
    # 还原
    bits_with_noise = bits_times_channel_coefficients_with_noise / h
    return bits_with_noise


def rician_torch(input_signal, snr_dB):
    K = 100
    N = input_signal.shape[0]
    device = input_signal.device

    # 生成 rician 信道复数项
    m = torch.randn(N, device=device)
    t = torch.randn(N, device=device)
    complex_mat = 1j * m + t    # (N, )
    # 计算 h
    h = (torch.sqrt(torch.tensor(K / (K + 1), device=device)) + torch.sqrt(torch.tensor(1 / (K +1), device=device))) \
        * complex_mat / torch.sqrt(torch.tensor(2., device=device))
    # 信道作用
    s = input_signal * h
    # 加噪声
    r = awgn_torch(s, snr_dB)
    bits_with_noise = r / h
    return bits_with_noise


class Channel(nn.Module):
    """
    Currently the channel model is either error free, erasure channel,
    rayleigh channel or the AWGN channel.
    """
    def __init__(self, channel_type, device=None):
        super(Channel, self).__init__()
        self.chan_type = channel_type
        self.device = device if device is not None else "cpu"

    def forward(self, x, snr):
        if self.chan_type == "awgn":
            return awgn_torch(x, snr)
        elif self.chan_type == "rayleigh":
            return rayleigh_torch(x, snr)
        elif self.chan_type == "rician":
            return rician_torch(x, snr)
        else:
            NotImplementedError("Unknown channel type " + self.chan_type)
